#!/bin/python
# -*- coding:utf-8 -*-
"""

"""
import sys
sys.path.append('..')

import gdal
import utils.raster as raster
import utils.text as text
import utils.common as common
import utils.stat_test as test
import numpy as np
import re
import os
import datetime
import math
import progressbar
import calendar
import random
import cv2
def get_lc_map(lc_name, codes=[1, 2, 3, 4, 5, 6, 7]):
	"""
	Args:

	Returns: a 0-1 ndarray

	Raises:
		RuntimeError:
	"""
	(lc_map_info, lc_map) = raster.open(lc_name)

	return lc_map_info, sum(lc_map == code for code in codes)

def get_vegetation_map(ndvi_name, thold=0.2):
	"""
	Args:

	Returns:

	Raises:
		RuntimeError:
	"""
	ndvi_info,ndvi_map = raster.open(ndvi_name)
	return (ndvi_map >= thold).astype(np.int16)

def burn_date_QC(burn_date,burn_date_qa,codes=[1,3]):
	"""
	Args:

	Returns:

	Raises:
		RuntimeError: 
	"""
	(bd_info,bd_map) = raster.open(burn_date)
	(bd_qc_info,bd_qc_map) = raster.open(burn_date_qa)

	return bd_info,bd_map * sum(bd_qc_map == qa_code for qa_code in codes)

	# return bd_info,bd_map * raster.detecte_isolate_pixel(sum(bd_qc_map == qa_code for qa_code in codes),window_size=3)

def extract_burn_event(root,out_file,var_compos=["ndvi","lai"],compos_step = 8,count = 4,var_daily = ["fmc","t2m","wns"],days=[0,1,2,3,4,5,6,7,14,21,28]):
	# 1. lc
	# 2. ndvi
	burn_date_file = common.list_file(os.path.join(root,"date"))
	## progress bar
	# total_progress = progressbar.ProgressBar()

	text.save(out_file,__gen_file_head__(var_compos,count,var_daily,days))

	dem = gdal.Open(os.path.join(root,"topo","dem.tif"))
	slope = gdal.Open(os.path.join(root,"topo","slope.tif"))
	aspect = gdal.Open(os.path.join(root,"topo","aspect.tif"))

	burt_event_seq = []
	with progressbar.ProgressBar(max_value = len(burn_date_file)) as total_progress:
		total_current = 0
		for bd_file in burn_date_file:
			print("Total:")
			total_current += 1
			total_progress.update(total_current)
			print('')
			print("Current Processing: %s"%bd_file)

			data_cube = {}
			date = re.sub("\D","",os.path.basename(bd_file))
			year = date[0:4]

			### 1. get the pixel calassfied as forest
			lc_year = __find_lc_nearest_year__(root,int(year))

			(lc_info,lc_map) = get_lc_map(os.path.join(root,"lc","lc_%d001.tif"%lc_year))
			(ld_info,ld_map) = get_lc_map(os.path.join(root,"lc","lc_%d001.tif"%lc_year),codes=[12,13,14])

			(bd_info,bd_map) = burn_date_QC(bd_file,os.path.join(root,"dqa","dqa_%s.tif"%date))

			# print(bd_map)

			burt_pixels = __burn_date_area__(bd_map)
			unburn_sample = __gen_random_point__([bd_info,bd_map],[lc_info,lc_map],15)
			# burt_pixels = np.where(bd_map > 0)
			
			
			data_cube = __construct_data_cube__(data_cube,root,date,var_compos+var_daily+["dlr","wnd"])
			
			# print(data_cube)


			with progressbar.ProgressBar(max_value = len(burt_pixels) + len(unburn_sample)) as current_progress:
				current_fire_id = 1
				for fire in burt_pixels:
					if len(burt_pixels[fire][0]) <= 1:
						current_progress.update(current_fire_id)
						current_fire_id += 1
						continue
					else:
						fire_pixels = burt_pixels[fire][0]
						current_progress.update(current_fire_id)
						current_fire_id += 1
						fire_event = {}
						# print(fire_pixels)
						for burt_pixel_idx in range(len(fire_pixels)):
							# print(current_fire_id)
							burn_date = "%s%03d"%(year,bd_map[fire_pixels[burt_pixel_idx][0]][fire_pixels[burt_pixel_idx][1]])

							burt_pixel_lat_lon = raster.imgXY2latlon([fire_pixels[burt_pixel_idx][1],fire_pixels[burt_pixel_idx][0]],bd_info[1])

							lc_map_XY = raster.latlont2imgXY(burt_pixel_lat_lon,lc_info[1])

							if not lc_map[math.floor(lc_map_XY[1])][math.floor(lc_map_XY[0])] == 1:
								continue
							else:
								fire_event.setdefault("lat",[]).append(burt_pixel_lat_lon[0])
								fire_event.setdefault("lon",[]).append(burt_pixel_lat_lon[1])
								fire_event.setdefault("alt",[]).append(__extract_value__(burt_pixel_lat_lon,dem))
								fire_event.setdefault("slope",[]).append(__extract_value__(burt_pixel_lat_lon,slope))
								fire_event.setdefault("aspect",[]).append(__extract_value__(burt_pixel_lat_lon,aspect))
								fire_event.setdefault("ld",[]).append(__get_nearest_distance__(ld_map,lc_map_XY)*500)
								fire_event.setdefault("dlr",[]).append(__extract_value__(burt_pixel_lat_lon,data_cube["dlr"][int(__get_doy__(burn_date,0))]))
								fire_event.setdefault("wnd",[]).append(__extract_value__(burt_pixel_lat_lon,data_cube["wnd"][int(__get_doy__(burn_date,0))]))

								for var in var_compos:
									nearest_date = "%s%03d"%(year,math.floor(int(burn_date[4:])/compos_step)*compos_step+1)

									for time in range(0,count):
										# print(__extract_value__(burt_pixel_lat_lon,data_cube[var][int(__get_doy__(nearest_date,-(time*compos_step)))]))
										# print(int(__get_doy__(nearest_date,-(time*compos_step),days = compos_step)))

										doy = __get_doy__(nearest_date,-(time*compos_step),days = compos_step)

										if int(doy) in data_cube[var]:
											fire_event.setdefault("%s_%d"%(var,time),[]).append(__extract_value__(burt_pixel_lat_lon,data_cube[var][int(doy)]))
										# event_seq_tmp = "%s\t%f" %(event_seq_tmp, __extract_value__(burt_pixel_lat_lon,data_cube[var][int(__get_doy__(nearest_date,-(time*compos_step)))]))

								for var in var_daily:
									var_data = []
									mul_data = __extract_value__(burt_pixel_lat_lon,data_cube["mul"][var])
									std_data = __extract_value__(burt_pixel_lat_lon,data_cube["std"][var])
									for time in days:
										doy = __get_doy__(burn_date,-time)
										if int(doy) in data_cube[var]:											
											var_data.append(__extract_value__(burt_pixel_lat_lon,data_cube[var][int(doy)]))

									# fire_event.setdefault("%s_%d"%(var,time),[]).append(__extract_value__(burt_pixel_lat_lon,data_cube[var][int(__get_doy__(burn_date,-time))]))

									# ascending by date
									var_data.reverse()
									(trend,z) = test.mk_trend_test(var_data)

									# print(var_data)
									# print(trend)

									if var is "tp":
										fire_event.setdefault(var,[]).append([np.sum(var_data),trend,z, 0 if std_data == 0 else (np.sum(var_data) - mul_data) / std_data])
									else:
										fire_event.setdefault(var,[]).append([np.mean(var_data),trend,z, 0 if std_data == 0 else (np.mean(var_data)- mul_data) / std_data])

									# fire_event.setdefault(var,{}).setdefault("a",(fire_event[var]["m"] - mul_data) / std_data)
										# event_seq_tmp = "%s\t%f" %(event_seq_tmp, __extract_value__(burt_pixel_lat_lon,data_cube[var][int(__get_doy__(burn_date,-time))]))
								# burt_event_seq.append("%s\n"%event_seq_tmp)	

						if len(fire_event) > 0:
							event_seq_tmp=datetime.datetime.strftime(datetime.datetime.strptime("%s%03d" %(year,burt_pixels[fire][1]),"%Y%j"),"%Y-%m-%d")

							for key in fire_event:
								# print(np.mean(fire_event[key]))
								if key in var_daily:
									# print()
									event_seq_tmp = "%s\t%s"%(event_seq_tmp,"\t".join(['{:f}'.format(x) for x in np.squeeze(np.mean(np.dstack(fire_event[key]),axis = 2)).tolist()]))

								else:
									event_seq_tmp = "%s\t%f"%(event_seq_tmp,np.mean(fire_event[key]))

							burt_event_seq.append("%s\tF1\n"%event_seq_tmp)
				for point in unburn_sample:
					current_progress.update(current_fire_id)
					current_fire_id += 1
					fire_event = {}
					burn_date = datetime.datetime.strptime(date,"%Y%j")

					burn_date = (burn_date + datetime.timedelta(days = calendar.monthrange(burn_date.year,burn_date.month)[1] - 1)).strftime("%Y%j")

					burt_pixel_lat_lon = raster.imgXY2latlon([point[1],point[0]],bd_info[1])

					lc_map_XY = raster.latlont2imgXY(burt_pixel_lat_lon,lc_info[1])
					if not lc_map[math.floor(lc_map_XY[1])][math.floor(lc_map_XY[0])] == 1:
						continue
					else:
						fire_event.setdefault("lat",[]).append(burt_pixel_lat_lon[0])
						fire_event.setdefault("lon",[]).append(burt_pixel_lat_lon[1])
						fire_event.setdefault("alt",[]).append(__extract_value__(burt_pixel_lat_lon,dem))
						fire_event.setdefault("slope",[]).append(__extract_value__(burt_pixel_lat_lon,slope))
						fire_event.setdefault("aspect",[]).append(__extract_value__(burt_pixel_lat_lon,aspect))
						fire_event.setdefault("ld",[]).append(__get_nearest_distance__(ld_map,lc_map_XY)*500)
						fire_event.setdefault("dlr",[]).append(__extract_value__(burt_pixel_lat_lon,data_cube["dlr"][int(__get_doy__(burn_date,0))]))
						fire_event.setdefault("wnd",[]).append(__extract_value__(burt_pixel_lat_lon,data_cube["wnd"][int(__get_doy__(burn_date,0))]))

						for var in var_compos:
							nearest_date = "%s%03d"%(year,math.floor(int(burn_date[4:])/compos_step)*compos_step+1)

							for time in range(0,count):
								doy = __get_doy__(nearest_date,-(time*compos_step),days = compos_step)

								if int(doy) in data_cube[var]:
									fire_event.setdefault("%s_%d"%(var,time),[]).append(__extract_value__(burt_pixel_lat_lon,data_cube[var][int(doy)]))
										

						for var in var_daily:
							var_data = []
							mul_data = __extract_value__(burt_pixel_lat_lon,data_cube["mul"][var])
							std_data = __extract_value__(burt_pixel_lat_lon,data_cube["std"][var])
							for time in days:
								doy = __get_doy__(burn_date,-time)
								if int(doy) in data_cube[var]:											
									var_data.append(__extract_value__(burt_pixel_lat_lon,data_cube[var][int(doy)]))

							# ascending by date
							var_data.reverse()
							(trend,z) = test.mk_trend_test(var_data)

							if var is "tp":
								fire_event.setdefault(var,[]).append([np.sum(var_data),trend,z, 0 if std_data == 0 else (np.sum(var_data) - mul_data) / std_data])
							else:
								fire_event.setdefault(var,[]).append([np.mean(var_data),trend,z, 0 if std_data == 0 else (np.mean(var_data)- mul_data) / std_data])

					if len(fire_event) > 0:
						event_seq_tmp=burn_date
						for key in fire_event:
									# print(np.mean(fire_event[key]))
							if key in var_daily:
										# print()
								event_seq_tmp = "%s\t%s"%(event_seq_tmp,"\t".join(['{:f}'.format(x) for x in np.squeeze(np.mean(np.dstack(fire_event[key]),axis = 2)).tolist()]))

							else:
								event_seq_tmp = "%s\t%f"%(event_seq_tmp,np.mean(fire_event[key]))

					burt_event_seq.append("%s\tF2\n"%event_seq_tmp)


						# print(fire_event)	
		text.save(out_file,burt_event_seq)

def __gen_random_point__(image, lc, num):
	# dx = [0,1,1,1,0,-1,-1,-1]
	# dy = [-1,-1,0,1,1,1,0,-1]

	# result = {}
	result = []
	if not image[1].dtype == np.float:
		image[1] = image[1].astype(np.float)

	lc[1] = lc[1].astype(np.float)

	im_fk3 = cv2.filter2D(image[1],-1,np.ones((11,11),np.float),borderType = 0)
	im_lc = cv2.filter2D(lc[1],-1,np.ones((3,3),np.float),borderType = 0)

	points = list(np.where(im_lc == 9))

	for point in random.sample([(points[0][i],points[1][i]) for i in range(points[0].shape[0])],num * 3):
		latlon = raster.imgXY2latlon([point[1],point[0]],lc[0][1])
		bd_xy = raster.latlont2imgXY(latlon,image[0][1])

		if im_fk3[math.floor(bd_xy[1])][math.floor(bd_xy[0])] == 0 :
			result.append((bd_xy[1],bd_xy[0]))

	if num >= len(result):
		return result
	else:
		return random.sample(result,num)

def __get_nearest_distance__(image,imgXY):
	imgXY = [math.floor(x) for x in imgXY]
	window_size = 1
	img_size = image.shape
	is_find = False
	distance = 0
	while True:
		for i in range(-window_size,window_size+1):
			# four direction
			if (imgXY[0] + window_size) < img_size[1] and (imgXY[1] + i) > 0 and (imgXY[1] + i) < img_size[0] and (image[imgXY[1] + i][imgXY[0]+window_size] == 1):
				distance = __euclidean_distance__([imgXY[0] + window_size,imgXY[1] + i],imgXY)
				is_find = True
				break
			if (imgXY[0] - window_size) > 0 and (imgXY[1] + i) > 0 and (imgXY[1] + i) < img_size[0] and (image[imgXY[1] + i][imgXY[0] - window_size] == 1):
				distance = __euclidean_distance__([imgXY[0] + window_size,imgXY[1] + i],imgXY)
				is_find = True
				break
			if (imgXY[1] + window_size < img_size[0]) and (imgXY[0] + i) > 0 and (imgXY[0] + i) < img_size[1] and (image[imgXY[1] + window_size][imgXY[0] + i] == 1):
				distance = __euclidean_distance__([imgXY[0] + window_size,imgXY[1] + i],imgXY)
				is_find = True
				break
			if (imgXY[1] - window_size > 0) and (imgXY[0] + i) > 0 and (imgXY[0] + i) < img_size[1] and (image[imgXY[1] - window_size][imgXY[0] + i] == 1):
				distance = __euclidean_distance__([imgXY[0] + window_size,imgXY[1] + i],imgXY)
				is_find = True
				break
		window_size += 1
		if is_find:
			break
	return distance

def __euclidean_distance__(src,dst):
	return math.sqrt(math.pow(src[0]-dst[0],2)+math.pow(src[1]-dst[1],2))

def __construct_data_cube__(data_cube,root,date,vars,backward_days = 62):

	data_cube_new = {}
	current_date = datetime.datetime.strptime(date,"%Y%j")


	last_day_of_month = current_date + datetime.timedelta(days = calendar.monthrange(current_date.year,current_date.month)[1])

	for day in range(backward_days):
		current_doy = (last_day_of_month + datetime.timedelta(days = -day)).strftime("%Y%j")

		for var in vars:
			if var in data_cube:
				if current_doy in data_cube[var]:
					data_cube_new.setdefault(var,{}).setdefault(current_doy,data_cube[var][current_doy])
				elif os.path.exists(os.path.join(root,var,"%s_%s.tif"%(var,current_doy))):
					data_cube_new.setdefault(var,{}).setdefault(int(current_doy),gdal.Open(os.path.join(root,var,"%s_%s.tif"%(var,current_doy))))
			elif os.path.exists(os.path.join(root,var,"%s_%s.tif"%(var,current_doy))):
					data_cube_new.setdefault(var,{}).setdefault(int(current_doy),gdal.Open(os.path.join(root,var,"%s_%s.tif"%(var,current_doy))))


	for var in vars:
		if os.path.exists(os.path.join(root,"monthly","mul_%s_%02d.tif"%(var,current_date.month))):
			data_cube_new.setdefault("mul",{}).setdefault(var,gdal.Open(os.path.join(root,"monthly","mul_%s_%02d.tif"%(var,current_date.month))))
			data_cube_new.setdefault("std",{}).setdefault(var,gdal.Open(os.path.join(root,"monthly","std_%s_%02d.tif"%(var,current_date.month))))

	return data_cube_new

def __gen_file_head__(var_compos,count,var_daily,days):
	tmp = "date\tlon\tlat\talt\tslope\taspect\tld\tdlr\twnd"
	for var in var_compos:
		for i in range(0, count):
			tmp = "%s\t%s_%d"%(tmp,var,i)

	for var in var_daily:
		tmp = "%s\t%s_m\t%s_t\t%s_z\t%s_a"%(tmp,var,var,var,var)
		# for i in days:
		# 	tmp = "%s\t%s_%d"%(tmp,var,i)

	return "%s\tfire\n"%tmp

def __burn_date_area__(image):

	areas_idx = {}

	dx = [0,1,1,1,0,-1,-1,-1]
	dy = [-1,-1,0,1,1,1,0,-1]

	idx = np.where(image > 0)
	xy_idx = [(idx[0][i],idx[1][i]) for i in range(idx[0].shape[0])]

	while len(xy_idx) > 0:
		x_y = [xy_idx.pop()]
		current_idx = [x_y[0]]
		date = [image[x_y[0][0]][x_y[0][1]]]

		if date[0] <= 0:
			continue

		visit = [x_y[0]]
		
		while len(x_y) > 0:
			x_y_current = x_y.pop()

			for i in range(0,8):
				
				nx = x_y_current[0] + dx[i]
				ny = x_y_current[1] + dy[i]

				if (nx,ny) in visit:
					continue
				else:
					visit.append((nx,ny))
				
				if nx >= 0 and nx < image.shape[0] and ny >= 0 and ny < image.shape[1]:
					if image[nx][ny] > 0 and abs(image[nx][ny] - image[x_y_current[0]][x_y_current[1]]) < 8:
						date.append(image[nx][ny])
						x_y.append((nx,ny))
						current_idx.append((nx,ny))

						if (nx,ny) in xy_idx:
							xy_idx.remove((nx,ny))

		areas_idx.setdefault(len(areas_idx),[current_idx,math.floor(np.mean(date))])
	return areas_idx

def __extract_value__(latlon,tiff_path):

	if isinstance(tiff_path,str):
		if os.path.exists(tiff_path):
			# (tiff_info,tiff_array) = raster.open(tiff_path)
			data = gdal.Open(tiff_path)
		else:
			return np.nan
	else:
		data = tiff_path

	tiff_info = raster.get_raster_info(data)

	tiff_XY = raster.latlont2imgXY(latlon,tiff_info[1])

	return data.GetRasterBand(1).ReadAsArray(math.floor(tiff_XY[0]),math.floor(tiff_XY[1]),1,1)[0][0]


def __get_doy__(date,delta,days = 1):
	"""
	Args:

	Returns:

	Raises:
		RuntimeError: 
	"""	
	init_date = datetime.datetime.strptime(date,"%Y%j")
	current_date = init_date + datetime.timedelta(days = delta)

	if current_date.year < init_date.year and days > 1:
		return "%d%03d"%(current_date.year,math.ceil(int(current_date.strftime("%j"))/days)*days+1)

	return current_date.strftime("%Y%j")

def __find_lc_nearest_year__(root,year):
	"""
	Args:

	Returns:

	Raises:
		RuntimeError: 
	"""
	while year > 1980:
		if os.path.exists(os.path.join(root,"lc","lc_%d001.tif"%year)):
			break
		else:
			year -= 1

	return year

if __name__ == '__main__':
	extract_burn_event(r"/mnt/fire/Extract","/mnt/fire/fire_event4.txt",var_daily=["fmc","rh","tp","t2m","wns"], days = [x for x in range(30)])
	# info,data = raster.open("/mnt/fire/Extract/date/date_2013001.tif")
	# __burn_date_area__(data)

